from fastapi import APIRouter, Query, HTTPException
from debate_orm import Team, RoundTeam, Score
from fastapi import WebSocket, WebSocketDisconnect
import uuid
from typing import Dict
app_judge = APIRouter()


"""
辅助函数块
"""


async def get_teams_in_contest(contest_id: int = Query(...), round_id: int = Query(...)):
    """
    获取某个赛次中某轮次比赛的所有参赛队伍
    """
    try:
        # 修正查询逻辑
        round_teams = await RoundTeam.filter(
            round__contest_id=contest_id,
            round__round_id=round_id
        ).prefetch_related('team')

        teams = [rt.team.team_id for rt in round_teams]
        return teams
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error fetching teams: {str(e)}")


# 存储活跃的 WebSocket 连接
active_connections: Dict[str, WebSocket] = {}


@app_judge.websocket("/judge/ws/scoring/{contest_id}/{round_id}")
async def websocket_scoring(websocket: WebSocket, contest_id: int, round_id: int):
    await websocket.accept()
    connection_id = str(uuid.uuid4())
    active_connections[connection_id] = websocket

    try:
        # 获取队伍列表
        team_list = await get_teams_in_contest(contest_id, round_id)
        current_index = 0

        # 存储所有队伍的分数（每个队伍一个分数列表）
        all_scores = {team_id: [] for team_id in team_list}

        # 发送第一个队伍信息
        if team_list:
            team_id = team_list[current_index]
            await websocket.send_json({
                "type": "request_score",
                "team_id": team_id,
                "current_index": current_index,
                "total_teams": len(team_list),
                "message": f"请为队伍 {team_id} 输入分数（可输入多个分数，用逗号分隔）",
                "input_type": "multiple"  # 表示需要输入多个分数
            })
        else:
            await websocket.send_json({
                "type": "error",
                "message": "没有找到参赛队伍"
            })
            return

        while current_index < len(team_list):
            # 等待前端发送分数
            data = await websocket.receive_json()

            if data.get("type") == "submit_score":
                scores_input = data.get("scores")  # 接收字符串形式的多个分数
                team_id = team_list[current_index]

                # 验证分数输入
                if not scores_input:
                    await websocket.send_json({
                        "type": "error",
                        "message": "请输入分数"
                    })
                    continue

                try:
                    # 解析多个分数
                    score_list = [int(score.strip()) for score in scores_input.split(",")]

                    # 存储到该队伍的分数列表中
                    all_scores[team_id] = score_list

                    # 发送确认信息
                    await websocket.send_json({
                        "type": "score_submitted",
                        "team_id": team_id,
                        "scores": score_list,
                        "message": f"队伍 {team_id} 分数提交成功: {score_list}"
                    })

                    # 移动到下一个队伍
                    current_index += 1

                    if current_index < len(team_list):
                        next_team_id = team_list[current_index]
                        await websocket.send_json({
                            "type": "request_score",
                            "team_id": next_team_id,
                            "current_index": current_index,
                            "total_teams": len(team_list),
                            "message": f"请为队伍 {next_team_id} 输入分数（可输入多个分数，用逗号分隔）",
                            "input_type": "multiple"
                        })
                    else:
                        # 所有队伍完成，保存到数据库
                        for team_id, scores in all_scores.items():
                            if scores:  # 只保存有分数的队伍
                                team_instance = await Team.get(team_id=team_id)
                                round_team_instance = await RoundTeam.get(
                                    team=team_instance,
                                    round__contest_id=contest_id,
                                    round__round_id=round_id
                                )

                                score_instance, created = await Score.get_or_create(
                                    team=round_team_instance,
                                    defaults={"score": scores}  # 存储分数列表
                                )

                                if not created:
                                    score_instance.score = scores
                                    await score_instance.save()

                        await websocket.send_json({
                            "type": "completed",
                            "message": "所有队伍打分完成！",
                            "summary": {
                                "total_teams": len(team_list),
                                "scored_teams": len([s for s in all_scores.values() if s]),
                                "all_scores": all_scores
                            }
                        })
                        break

                except ValueError:
                    await websocket.send_json({
                        "type": "error",
                        "message": "请输入有效的数字分数，多个分数用逗号分隔"
                    })
                except Team.DoesNotExist:
                    await websocket.send_json({
                        "type": "error",
                        "message": f"队伍 {team_id} 不存在"
                    })
                except RoundTeam.DoesNotExist:
                    await websocket.send_json({
                        "type": "error",
                        "message": f"队伍 {team_id} 未参加本轮比赛"
                    })

                    await websocket.send_json({
                        "type": "completed",
                        "message": "所有队伍处理完成",
                        "summary": {
                            "total_teams": len(team_list),
                            "scored_teams": len([s for s in all_scores.values() if s]),
                            "all_scores": all_scores
                        }
                    })
                    break

    except WebSocketDisconnect:
        print(f"客户端断开连接: {connection_id}")
    except Exception as e:
        await websocket.send_json({
            "type": "error",
            "message": f"系统错误: {str(e)}"
        })
    finally:
        # 清理连接
        if connection_id in active_connections:
            del active_connections[connection_id]
